import chromadb
import os

def setup_chroma_database(db_path="./memory_db"):
    """
    设置Chroma数据库
    """
    try:
        # 创建持久化客户端
        client = chromadb.PersistentClient(path=db_path)
        
        # 创建医学记忆集合
        collection = client.get_or_create_collection(
            name="medical_memories",
            metadata={"description": "智能医学工程虚拟助手记忆库"}
        )
        
        print(f"✅ Chroma数据库已设置，路径: {os.path.abspath(db_path)}")
        return client, collection
        
    except Exception as e:
        print(f"❌ 数据库设置失败: {e}")
        return None, None

if __name__ == "__main__":
    client, collection = setup_chroma_database()
    if client and collection:
        print("数据库初始化成功！")